#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
键盘检测模块 - 包含高级图像处理和按键识别算法
"""

import cv2
import numpy as np
import numpy.typing as npt
from typing import Tuple, List, Dict, Optional, Any
import logging
from dataclasses import dataclass
from enum import Enum

class KeyIssueType(Enum):
    MISSING = "缺失"
    MISPLACED = "位置错误"
    ROTATED = "方向错误"
    WRONG_KEY = "错误按键"

@dataclass
class KeyIssue:
    issue_type: KeyIssueType
    position: Tuple[int, int]
    description: str
    confidence: float

class KeyboardDetector:
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        
        # 键盘检测参数
        self.min_keyboard_area = 50000  # 最小键盘区域面积
        self.key_contour_area_range = (100, 5000)  # 按键轮廓面积范围
        self.aspect_ratio_tolerance = 0.3  # 长宽比容差
        
        # 模板匹配参数
        self.match_threshold = 0.7  # 匹配阈值
        self.rotation_angles = [0, 90, 180, 270]  # 检测旋转角度
        
    def detect_keyboard_keys(self, image: npt.NDArray) -> List[Dict[str, Any]]:
        """
        检测键盘上的所有按键
        
        Args:
            image: 输入图像
            
        Returns:
            按键信息列表，每个按键包含位置、大小、角度等信息
        """
        # 预处理图像
        processed = self._preprocess_for_key_detection(image)
        
        # 检测按键轮廓
        key_contours = self._detect_key_contours(processed)
        
        # 提取按键信息
        keys = []
        for contour in key_contours:
            key_info = self._extract_key_info(contour, image)
            if key_info:
                keys.append(key_info)
                
        # 按位置排序（从上到下，从左到右）
        keys.sort(key=lambda k: (k['y'], k['x']))
        
        return keys
        
    def _preprocess_for_key_detection(self, image: npt.NDArray) -> npt.NDArray:
        """为按键检测预处理图像"""
        # 转换为灰度图像
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        else:
            gray = image.copy()
            
        # 自适应直方图均衡化
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        enhanced = clahe.apply(gray)
        
        # 高斯模糊
        blurred = cv2.GaussianBlur(enhanced, (5, 5), 0)
        
        # 自适应阈值处理
        binary = cv2.adaptiveThreshold(
            blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
            cv2.THRESH_BINARY_INV, 11, 2
        )
        
        # 形态学操作
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
        cleaned = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
        
        return cleaned
        
    def _detect_key_contours(self, binary_image: npt.NDArray) -> List[npt.NDArray]:
        """检测按键轮廓"""
        contours, _ = cv2.findContours(
            binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )
        
        # 过滤轮廓
        key_contours = []
        for contour in contours:
            area = cv2.contourArea(contour)
            if self.key_contour_area_range[0] <= area <= self.key_contour_area_range[1]:
                # 检查长宽比
                x, y, w, h = cv2.boundingRect(contour)
                aspect_ratio = w / h if h > 0 else 0
                
                if 0.5 <= aspect_ratio <= 2.0:  # 合理的长宽比范围
                    key_contours.append(contour)
                    
        return key_contours
        
    def _extract_key_info(self, contour: npt.NDArray, original_image: npt.NDArray) -> Optional[Dict[str, Any]]:
        """提取按键信息"""
        # 获取边界框
        x, y, w, h = cv2.boundingRect(contour)
        
        # 计算最小外接矩形
        rect = cv2.minAreaRect(contour)
        center, size, angle = rect
        
        # 提取按键区域
        key_region = original_image[y:y+h, x:x+w]
        
        # 计算按键特征
        key_info = {
            'x': int(center[0]),
            'y': int(center[1]),
            'width': int(size[0]),
            'height': int(size[1]),
            'angle': angle,
            'area': cv2.contourArea(contour),
            'region': key_region,
            'contour': contour
        }
        
        return key_info
        
    def compare_keyboards(self, reference_keys: List[Dict], current_keys: List[Dict]) -> List[KeyIssue]:
        """
        对比两个键盘的按键配置
        
        Args:
            reference_keys: 参考键盘按键列表
            current_keys: 当前键盘按键列表
            
        Returns:
            检测到的问题列表
        """
        issues = []
        
        # 检查按键数量
        if len(current_keys) < len(reference_keys):
            issues.append(KeyIssue(
                issue_type=KeyIssueType.MISSING,
                position=(0, 0),
                description=f"检测到 {len(reference_keys) - len(current_keys)} 个按键缺失",
                confidence=0.9
            ))
            
        # 为每个参考按键找到最近的当前按键
        used_current_indices = set()
        
        for ref_key in reference_keys:
            best_match_idx = None
            best_match_distance = float('inf')
            best_match_score = 0
            
            for i, curr_key in enumerate(current_keys):
                if i in used_current_indices:
                    continue
                    
                # 计算位置距离
                distance = np.sqrt(
                    (ref_key['x'] - curr_key['x'])**2 + 
                    (ref_key['y'] - curr_key['y'])**2
                )
                
                # 计算相似度分数
                similarity = self._calculate_key_similarity(ref_key, curr_key)
                
                # 综合评分
                if distance < 50 and similarity > self.match_threshold:
                    if distance < best_match_distance:
                        best_match_idx = i
                        best_match_distance = distance
                        best_match_score = similarity
                        
            if best_match_idx is not None:
                used_current_indices.add(best_match_idx)
                curr_key = current_keys[best_match_idx]
                
                # 检查具体问题
                issues.extend(self._check_key_issues(ref_key, curr_key))
            else:
                # 未找到匹配的按键
                issues.append(KeyIssue(
                    issue_type=KeyIssueType.MISSING,
                    position=(ref_key['x'], ref_key['y']),
                    description=f"在位置 ({ref_key['x']}, {ref_key['y']}) 处缺失按键",
                    confidence=0.8
                ))
                
        return issues
        
    def _calculate_key_similarity(self, key1: Dict, key2: Dict) -> float:
        """计算两个按键的相似度"""
        # 位置相似度
        pos_distance = np.sqrt(
            (key1['x'] - key2['x'])**2 + (key1['y'] - key2['y'])**2
        )
        pos_similarity = max(0, 1 - pos_distance / 100)
        
        # 大小相似度
        size_diff = abs(key1['width'] - key2['width']) + abs(key1['height'] - key2['height'])
        size_similarity = max(0, 1 - size_diff / 100)
        
        # 角度相似度
        angle_diff = abs(key1['angle'] - key2['angle']) % 360
        if angle_diff > 180:
            angle_diff = 360 - angle_diff
        angle_similarity = max(0, 1 - angle_diff / 90)
        
        # 综合相似度
        total_similarity = (pos_similarity * 0.4 + size_similarity * 0.3 + angle_similarity * 0.3)
        
        return total_similarity
        
    def _check_key_issues(self, ref_key: Dict, curr_key: Dict) -> List[KeyIssue]:
        """检查按键的具体问题"""
        issues = []
        
        # 检查位置偏差
        pos_distance = np.sqrt(
            (ref_key['x'] - curr_key['x'])**2 + (ref_key['y'] - curr_key['y'])**2
        )
        
        if pos_distance > 20:  # 位置偏差阈值
            issues.append(KeyIssue(
                issue_type=KeyIssueType.MISPLACED,
                position=(curr_key['x'], curr_key['y']),
                description=f"按键位置偏差 {pos_distance:.1f} 像素",
                confidence=min(0.9, pos_distance / 50)
            ))
            
        # 检查旋转
        angle_diff = abs(ref_key['angle'] - curr_key['angle']) % 360
        if angle_diff > 180:
            angle_diff = 360 - angle_diff
            
        if angle_diff > 10:  # 旋转偏差阈值
            issues.append(KeyIssue(
                issue_type=KeyIssueType.ROTATED,
                position=(curr_key['x'], curr_key['y']),
                description=f"按键旋转 {angle_diff:.1f} 度",
                confidence=min(0.9, angle_diff / 45)
            ))
            
        # 检查大小差异
        width_diff = abs(ref_key['width'] - curr_key['width']) / ref_key['width']
        height_diff = abs(ref_key['height'] - curr_key['height']) / ref_key['height']
        
        if width_diff > 0.2 or height_diff > 0.2:
            issues.append(KeyIssue(
                issue_type=KeyIssueType.WRONG_KEY,
                position=(curr_key['x'], curr_key['y']),
                description=f"按键大小异常 (宽差: {width_diff:.1%}, 高差: {height_diff:.1%})",
                confidence=max(width_diff, height_diff)
            ))
            
        return issues
        
    def visualize_detection(self, image: npt.NDArray, keys: List[Dict], issues: List[KeyIssue]) -> npt.NDArray:
        """可视化检测结果"""
        vis_image = image.copy()
        
        # 绘制检测到的按键
        for key in keys:
            # 绘制按键边界框
            cv2.rectangle(
                vis_image,
                (key['x'] - key['width']//2, key['y'] - key['height']//2),
                (key['x'] + key['width']//2, key['y'] + key['height']//2),
                (0, 255, 0), 2
            )
            
            # 绘制按键中心点
            cv2.circle(vis_image, (key['x'], key['y']), 3, (0, 0, 255), -1)
            
        # 绘制问题标记
        for issue in issues:
            x, y = issue.position
            
            if issue.issue_type == KeyIssueType.MISSING:
                color = (255, 0, 0)  # 蓝色 - 缺失
                symbol = "X"
            elif issue.issue_type == KeyIssueType.MISPLACED:
                color = (0, 255, 255)  # 黄色 - 位置错误
                symbol = "M"
            elif issue.issue_type == KeyIssueType.ROTATED:
                color = (255, 255, 0)  # 青色 - 旋转
                symbol = "R"
            else:
                color = (255, 0, 255)  # 紫色 - 错误按键
                symbol = "W"
                
            # 绘制问题标记
            cv2.putText(
                vis_image, symbol, (x - 10, y - 10),
                cv2.FONT_HERSHEY_SIMPLEX, 1, color, 2
            )
            
            # 绘制问题圆圈
            cv2.circle(vis_image, (x, y), 15, color, 2)
            
        # 添加统计信息
        total_keys = len(keys)
        total_issues = len(issues)
        
        info_text = f"检测到按键: {total_keys}, 问题数量: {total_issues}"
        cv2.putText(
            vis_image, info_text, (10, 30),
            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2
        )
        
        return vis_image
        
    def save_detection_results(self, image_path: str, keys: List[Dict], issues: List[KeyIssue]):
        """保存检测结果到文件"""
        results = {
            'timestamp': cv2.getTickCount(),
            'total_keys': len(keys),
            'total_issues': len(issues),
            'keys': keys,
            'issues': [
                {
                    'type': issue.issue_type.value,
                    'position': issue.position,
                    'description': issue.description,
                    'confidence': issue.confidence
                }
                for issue in issues
            ]
        }
        
        # 保存为JSON文件
        import json
        with open(image_path + '_detection_results.json', 'w', encoding='utf-8') as f:
            json.dump(results, f, ensure_ascii=False, indent=2)
            
    def load_reference_template(self, template_path: str) -> Dict:
        """加载参考键盘模板"""
        try:
            # 读取模板图像
            template_image = cv2.imread(template_path)
            if template_image is None:
                raise ValueError(f"无法加载模板图像: {template_path}")
                
            # 检测按键
            keys = self.detect_keyboard_keys(template_image)
            
            # 创建模板信息
            template_info = {
                'image': template_image,
                'keys': keys,
                'key_count': len(keys),
                'image_path': template_path
            }
            
            return template_info
            
        except Exception as e:
            self.logger.error(f"加载模板失败: {str(e)}")
            raise